from typing import *
from heapq import *


class Solution:
    def findMaximumElegance(self, items: List[List[int]], k: int) -> int:
        # items[i] = [profiti, categoryi]
        items.sort(key=lambda k: k[0], reverse=True)
        cnt = Counter()
        sm = 0
        q = []
        for i in range(0, k):
            p, c = items[i]
            cnt[c] += 1
            sm += p
            if cnt[c] > 1:
                heappush(q, (p, c))
        ans = sm + len(cnt) * len(cnt)
        for i in range(k, len(items)):
            p, c = items[i]
            if c not in cnt and len(q) > 0:
                sm -= q[0][0] - p
                heappop(q)
                cnt[c] += 1
            ans = max(ans, sm + len(cnt) * len(cnt))
        return ans


s = Solution()
print(
    s.findMaximumElegance(
        [
            [10, 1],
            [10, 1],
            [10, 1],
            [10, 1],
            [10, 1],
            [10, 1],
            [10, 1],
            [10, 1],
            [10, 1],
            [10, 1],
            [3, 2],
            [3, 3],
            [3, 4],
            [3, 5],
            [3, 6],
            [3, 7],
            [3, 8],
            [3, 9],
            [3, 10],
            [3, 11],
        ],
        10,
    )
)

# print(
#     s.findMaximumElegance(
#         [[1, 5], [4, 5], [9, 8], [6, 8], [4, 7], [4, 2], [4, 3], [10, 8]], 4
#     )
# )
# print(s.findMaximumElegance([[3, 4], [8, 4], [2, 2], [1, 3]], 2))
# print(s.findMaximumElegance([[5, 1], [6, 1], [8, 1]], 2))
# print(s.findMaximumElegance([[1, 2], [6, 2]], 1))
# print(s.findMaximumElegance(items=[[3, 2], [5, 1], [10, 1]], k=2))

# cnt = Counter()
# cnt[1] += 1
# print(len(cnt), cnt[1])
# cnt[1] -= 1
# if cnt[1] == 0:
#     cnt.pop(1)
# print(len(cnt), cnt[1])
